# AI画图
from web.apis.TengXunYuanBao import JiChuBan
import requests
import json


class TengXunYuanBao_GenerateImages(JiChuBan.YuanBaoJiChuBan):
    def __init__(self):
        super().__init__()

    # 生成图片
    def generate_images(self, prompt, imageIntention, chat_id):
        data = {"model":"gpt_175B_0404","prompt":prompt,"plugin":"Adaptive","displayPrompt":f"{prompt} {imageIntention}风格 ","displayPromptType":1,"options":{"imageIntention":{"needIntentionModel":True,"backendUpdateFlag":2,"userIntention":{"style":imageIntention + "风格"},"intentionStatus":True}},"displayImageIntentionLabels":[{"type":"style","disPlayValue":imageIntention+"风格","startIndex":0,"endIndex":5}],"multimedia":[],"agentId":self.GenerateImages_id,"supportHint":2,"version":"v2"}
        data = json.dumps(data)
        response = requests.post(f'{self.domain}/api/chat/{chat_id}',
                                 cookies=self.cookies, headers=self.headers, data=data, stream=True)
        # self.handle_stream_response(response)
        return response
    
    # 获取风格
    def get_style(self):
        json_data = {
            'agentId': self.GenerateImages_id,
        }

        response = requests.post('https://yuanbao.tencent.com/api/user/agent/detail', cookies=self.cookies, headers=self.headers, json=json_data)
        return response.json()



if __name__ == "__main__":
    # TengXunYuanBao_GenerateImages().generate_images("狮子", "赛博朋克风格", "a47eda39-7c12-4e97-99d6-057d12526dea")
    yuanbao = TengXunYuanBao_GenerateImages()
    print(yuanbao.get_style())
